#include <ros/ros.h>
#include <geometry_msgs/PoseStamped.h>
#include <mavros_msgs/CommandBool.h>
#include <mavros_msgs/SetMode.h>
#include <mavros_msgs/State.h>
#include <cmath>

mavros_msgs::State current_state;
geometry_msgs::PoseStamped current_pose;

void state_cb(const mavros_msgs::State::ConstPtr& msg){
    current_state = *msg;
}

void pose_cb(const geometry_msgs::PoseStamped::ConstPtr& msg){
    current_pose = *msg;
}

int main(int argc, char **argv)
{
    ros::init(argc, argv, "offb_node");
    ros::NodeHandle nh;

    ros::Subscriber state_sub = nh.subscribe<mavros_msgs::State>
            ("mavros/state", 10, state_cb);
    ros::Subscriber pose_sub = nh.subscribe<geometry_msgs::PoseStamped>
            ("mavros/local_position/pose", 10, pose_cb);
    ros::Publisher local_pos_pub = nh.advertise<geometry_msgs::PoseStamped>
            ("mavros/setpoint_position/local", 10);
    ros::ServiceClient arming_client = nh.serviceClient<mavros_msgs::CommandBool>
            ("mavros/cmd/arming");
    ros::ServiceClient set_mode_client = nh.serviceClient<mavros_msgs::SetMode>
            ("mavros/set_mode");

    ros::Rate rate(20.0);

    while(ros::ok() && !current_state.connected){
        ros::spinOnce();
        rate.sleep();
    }

    // Define waypoints to form a square
    geometry_msgs::PoseStamped waypoints[5];
    waypoints[0].pose.position.x = 0;
    waypoints[0].pose.position.y = 0;
    waypoints[0].pose.position.z = 2;

    waypoints[1].pose.position.x = 2;
    waypoints[1].pose.position.y = 0;
    waypoints[1].pose.position.z = 2;

    waypoints[2].pose.position.x = 2;
    waypoints[2].pose.position.y = 2;
    waypoints[2].pose.position.z = 2;

    waypoints[3].pose.position.x = 0;
    waypoints[3].pose.position.y = 2;
    waypoints[3].pose.position.z = 2;

    waypoints[4].pose.position.x = 0;
    waypoints[4].pose.position.y = 0;
    waypoints[4].pose.position.z = 2;

    // Send initial setpoints
    for(int i = 100; ros::ok() && i > 0; --i){
        local_pos_pub.publish(waypoints[0]);
        ros::spinOnce();
        rate.sleep();
    }

    mavros_msgs::SetMode offb_set_mode;
    offb_set_mode.request.custom_mode = "OFFBOARD";

    mavros_msgs::CommandBool arm_cmd;
    arm_cmd.request.value = true;

    if( set_mode_client.call(offb_set_mode) && offb_set_mode.response.mode_sent){
        ROS_INFO("Offboard enabled");
    }
    if( arming_client.call(arm_cmd) && arm_cmd.response.success){
        ROS_INFO("Vehicle armed");
    }

    int waypoint_index = 0;

    while(ros::ok()){
        // Check if we are close to the current waypoint
        float dx = current_pose.pose.position.x - waypoints[waypoint_index].pose.position.x;
        float dy = current_pose.pose.position.y - waypoints[waypoint_index].pose.position.y;

        if (sqrt(dx*dx + dy*dy) < 0.1) {
            waypoint_index = (waypoint_index + 1) % 5;  // Move to next waypoint
        }

        local_pos_pub.publish(waypoints[waypoint_index]);

        ros::spinOnce();
        rate.sleep();
    }

    return 0;
}
